import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
import copy


class Agent():

    def __init__(self):
        self.time = 0.5
        self.mass = 80
        self.desiredVel = 0.8
        self.charactime = 0.5
        self.position = np.array([np.random.uniform(20, 100), np.random.uniform(20, 100)])
        self.destination = np.array([110, 60])
        self.actualVel = np.array([0, 0])
        self.radius = 0.3
        self.constantA = 2000
        self.constantB = 0.08
        self.parambody = 120000
        self.paramsliding = 240000
        self.final = np.array([130, 60])

    def g(self, a, b):
        # g函数，在社会力的计算中会用到
        return 0 if a < b else a - b

    def distanceP2W(self, point, wall):
        # 计算点到直线的距离，返回[距离，方向单位向量]
        p0 = np.array(wall[0])
        p1 = np.array(wall[1])
        b = p1 - p0
        a = point - p0
        c_m = np.dot(a, b) / np.sqrt(np.dot(b, b))
        c = c_m * b / np.sqrt(np.dot(b, b))
        e = a - c
        dist = np.sqrt(np.dot(e, e))
        npw = e / dist
        return dist, npw

    def personalDesireForce(self):
        # 计算人的自驱动力
        direction = (self.position - self.destination) / np.linalg.norm(self.position - self.destination) if \
            self.position[0] < 109 else (self.position - self.final) / np.linalg.norm(self.position - self.final)
        return -self.mass * (self.desiredVel * direction - self.actualVel) / self.charactime

    def peopleInteraction(self, other):
        # 计算人与人之间的作用力
        rij = self.radius + other.radius
        dij = np.linalg.norm(self.position - other.position)
        nij = (self.position - other.position) / dij
        tij = np.array([-nij[1], nij[0]])
        tangVelDiff = (self.actualVel - other.actualVel) * tij
        return (self.constantA * np.exp((rij - dij) / self.constantB) + self.parambody * self.g(rij, dij)) * nij - self.paramsliding * self.g(
            rij, dij) * tangVelDiff * tij

    def wallInteraction(self, wall):
        # 计算墙体对人的作用力
        diw, niw = self.distanceP2W(self.position, wall)
        tiw = np.array([-niw[1], niw[0]])
        return (self.constantA * np.exp((self.radius - diw) / self.constantB) + self.parambody * self.g(self.radius, diw)) * niw - self.paramsliding * self.g(
            self.radius, diw) * (self.actualVel * tiw) * tiw

    def changeFunc(self, Agents, Walls):
        # 计算总的受力，加速对，速度与移动后的位置
        sumForce = 0
        sumPeopleForce = 0
        for i in Agents:
            if (i.position == self.position).all():
                continue
            else:
                sumPeopleForce += self.peopleInteraction(i)
        sumWallForce = 0
        for i in Walls:
            if np.dot(i[0] - self.position, i[0] - i[1]) < 0 or np.dot(i[1] - self.position, i[1] - i[0]) < 0:
                pass
            else:
                sumWallForce += self.wallInteraction(i)
        sumForce = sumPeopleForce + sumWallForce + self.personalDesireForce()
        accl = sumForce / self.mass
        actualVel = accl * self.time + self.actualVel + sumWallForce
        self.actualVel = actualVel if np.linalg.norm(actualVel) < 1 else actualVel / np.linalg.norm(actualVel)
        self.position += self.actualVel * self.time


def Draw(Agents, Walls):
    plt.cla()
    plt.xlim((0, 120))
    plt.ylim((0, 120))
    plt.axis([0, 120, 0, 120])
    copyAgents = copy.deepcopy(Agents)
    newAgents = copy.deepcopy(Agents)
    for i in range(len(Agents)):
        newAgents[i].changeFunc(copyAgents, Walls)
        a = newAgents[i].position[0]
        b = newAgents[i].position[1]
        theta = np.arange(0, 2 * np.pi, 0.01)
        x = a + np.dot(Agents[i].radius, np.cos(theta));
        y = b + np.dot(Agents[i].radius, np.sin(theta))
        axes = fig.add_subplot(111);
        axes.plot(x, y, color='black')
    for i in Walls:
        plt.plot([i[0][0], i[1][0]], [i[0][1], i[1][1]], color='black')
    Agents[:] = newAgents[:]


# 调用animation.FuncAnimation模块
def init(Agents, Walls):
    return Draw(Agents, Walls)


def update(frameNum, Agents, Walls):
    for i in Agents:
        if i.position[0] > 130:
            Agents.remove(i)
    return Draw(Agents, Walls)


if __name__ == '__main__':

    AgeNum = 100
    Agents = []
    while len(Agents) < AgeNum:
        agent = Agent()
        flag = False
        for j in range(len(Agents)):
            if np.linalg.norm(agent.position - Agents[j].position) < (agent.radius + Agents[j].radius):
                flag = True
            if flag:
                break
        if flag:
            pass
        else:
            Agents.append(agent)

    Walls = [
        [np.array([10, 10]), np.array([110, 10]), 5],
        [np.array([10, 10]), np.array([10, 110]), 5],
        [np.array([10, 110]), np.array([110, 110]), 5],
        [np.array([110, 10]), np.array([110, 59]), 5],
        [np.array([110, 110]), np.array([110, 61]), 5]
    ]
    # 设置画布参数与画图函数
    fig = plt.figure()

    ani = animation.FuncAnimation(fig, update, frames=100, fargs=(Agents, Walls), interval=100)
    plt.show()
